import hashlib
import optparse
import os
import random
import string
import time
from typing import Dict, List
from urllib.parse import urlencode

from Crypto.Cipher import AES
from binascii import b2a_hex, a2b_hex
from Crypto.Util.Padding import pad, unpad


class KeyCheckError(Exception):
    pass


class KeyManager(object):
    """
    秘钥管理器

    所需环境:
    pip install py-dmidecode==0.1.0
    sudo apt install hwinfo

    脚本运行示例:
        验证key:
            python key_manager.py
            python key_manager.py -k {key}
        创建key:
            python key_manager.py -m
            python key_manager.py -m -t '2022-10-10 10:10:10'
    """
    key_file = './key_file.txt'  # 秘钥存放位置

    symmetric_key = 'hiU49YHf2Nv4Y5I7M6K3izvTWYtALha9'
    random_bytes = b"r\xefG\x02t\xc9\x91^\x8b$-\x90\xb0\xee\xb5\xa6"

    def __init__(self, key_file=None, symmetric_key=None):
        self.key_file = key_file or self.key_file
        self.symmetric_key = symmetric_key or self.symmetric_key

    @classmethod
    def symmetric_encrypt(cls, text: str = ''):
        """加密"""
        crypto = AES.new(cls.symmetric_key.encode("utf-8"), AES.MODE_OFB, cls.random_bytes)
        data = pad(text.encode("utf-8"), AES.block_size, style="x923")
        encrypt_data = b2a_hex(crypto.encrypt(data)).decode("utf-8")
        return encrypt_data

    @classmethod
    def symmetric_decrypt(cls, text: str = ''):
        """解密"""
        crypto = AES.new(cls.symmetric_key.encode("utf-8"), AES.MODE_OFB, cls.random_bytes)
        data = crypto.decrypt(a2b_hex(text))
        decrypt_data = unpad(data, AES.block_size, style="x923").decode()
        return decrypt_data

    def make_key(self, expiration: str = None):
        """
        生成秘钥
        key1-key2-key3  / key1-key2
        key1:
            硬盘序列号以'-'拼接完整字符串, 再进行对称加密
        key2:
            系统唯一标识 + 随机生成的数字字符串 进行md5加密, 最后拼接随机生成的数字字符串
        key3:
            过期时间(提供过期时间时生成)
        :param expiration: 过期时间 格式: %Y-%m-%d %H:%M:%S
        :return:
        """
        disk_ids = self.get_disk_serial_ids()
        system_info = self.get_system_info()

        len_disk = len(disk_ids)
        key1 = self.symmetric_encrypt('-'.join(disk_ids))
        salt = "".join(random.choices(string.digits, k=len_disk))

        if expiration is not None:
            # 包含过期时间
            expiration_time = int(time.mktime(time.strptime(expiration, '%Y-%m-%d %H:%M:%S')))
            key = f"{key1}-{self.salt_params(**system_info, salt=salt, expiration_time=expiration_time)}" \
                  f"{salt}-{expiration_time}"
        else:
            key = f"{key1}-{self.salt_params(**system_info, salt=salt)}{salt}"
        with open(self.key_file, 'w', encoding='utf-8') as f:
            f.write(key)

    def check_key(self, key=None) -> bool:
        """
        检查key是否有效
        False=验证不通过 True=通过
        """
        try:
            # 打开文件
            if not key:
                try:
                    with open(self.key_file, 'r', encoding='utf-8') as f:
                        key = ''.join(f.read().split())
                except Exception:
                    raise KeyCheckError('获取本地加密字符串失败!')

            disk_ids = self.get_disk_serial_ids()
            system_info = self.get_system_info()

            keys = key.split('-')
            key1, key2 = keys[0], keys[1]
            # key1 验证
            try:
                str_list = self.symmetric_decrypt(key1).split('-')
                disk_len = len(str_list)
                for disk_id in disk_ids:
                    if disk_id in str_list:
                        str_list.remove(disk_id)
            except Exception:
                raise KeyCheckError('校验失败!')
            if str_list:
                raise KeyCheckError('校验失败!')

            # 其他key 验证
            expiration_time = None
            try:
                key2 = keys[1]
                salt = key2[-disk_len:]
                if len(keys) == 3:
                    expiration_time = int(keys[2])
                    sign = self.salt_params(**system_info, salt=salt, expiration_time=expiration_time)
                else:
                    sign = self.salt_params(**system_info, salt=salt)
            except Exception:
                raise KeyCheckError('校验失败!')
            if f"{sign}{salt}" != key2:
                raise KeyCheckError('校验失败!')
            if expiration_time is not None and (expiration_time - int(time.time())) <= 0:
                raise KeyCheckError('验证码已过期!')
        except Exception as e:
            print(e)
            return False
        else:
            return True

    @staticmethod
    def get_system_info() -> Dict:
        """
        获取系统唯一信息
        """
        from dmidecode import DMIDecode

        dmi = DMIDecode()
        data = {}
        for k, v in dmi.data.items():
            if v.get('DMIName') == 'System Information':
                data = v
                break
        if not data:
            raise KeyCheckError('未获取到系统信息')
        return {'sys_serial_number': data['Serial Number'], 'sys_uuid': data['UUID']}

    # @staticmethod
    # def __hard_disk_uuids() -> List:
    #     """
    #     获取挂载硬盘的uuids
    #     """
    #     lines = os.popen("ls -lh /dev/disk/by-uuid/ | awk '{print $9}'").readlines()
    #     r = []
    #     for line in lines:
    #         line = line.strip()
    #         if line:
    #             r.append(line)
    #     return r

    @staticmethod
    def get_disk_serial_ids() -> List:
        """
        获取硬盘序列号
        :return:
        """
        lines = os.popen("hwinfo --disk | grep 'Serial ID' | awk '{print $3}'").readlines()
        r = []
        for line in lines:
            line = line.strip()
            line = line[1:-1]  # 去除可能的 "
            if line:
                r.append(line)
        if not r:
            raise KeyCheckError('未获取到硬盘信息')
        return r

    @staticmethod
    def salt_params(sys_serial_number=None, sys_uuid=None, salt=None, **kwargs):
        """
        加密参数
        """
        if not all([
            sys_serial_number,
            sys_uuid,
            salt
        ]):
            raise KeyCheckError('待加密数据异常')

        key = (
            hashlib.md5(
                f"andiSec,{sys_serial_number},{sys_uuid},{salt}".encode(encoding="UTF-8")
            )
        ).hexdigest()
        params = {
            'ssn': sys_serial_number,
            'uuid': sys_uuid,
            'salt': salt
        }
        params.update(kwargs)
        sign = (
            hashlib.md5(
                f"{key[:9]}{urlencode(params)}{key[9:]}".encode(encoding="UTF-8")
            )
        ).hexdigest()
        return sign


if __name__ == '__main__':
    usage = 'Usage: %prog -c/-a -t <time_str> -k <key_str> -f <key_file> --sk <symmetric_key>'
    parser = optparse.OptionParser(usage, version='v1.0')
    parser.add_option('-m', dest='make', action='store_true', help='创建秘钥', default=False)
    parser.add_option('-t', dest='expiration_str', type='string', help='过期时间 %Y-%m-%d %H:%M:%S')
    parser.add_option('-f', dest='key_file', type='string', help='秘钥位置')
    parser.add_option('-k', dest='key', type='string', help='秘钥')
    parser.add_option('--sk', dest='symmetric_key', type='string', help='对称加密字符(8位)')
    parser.add_option('--test', dest='test', action='store_true', help='测试', default=False)

    options, args = parser.parse_args()

    key_manager = KeyManager(key_file=options.key_file, symmetric_key=options.symmetric_key)

    if options.test:
        try:

            print(key_manager.get_system_info())
            print(key_manager.get_disk_serial_ids())
        except Exception as e:
            print(f'error >>>>>> {e}')
        else:
            print('>>>>>> success <<<<<<<')
    else:
        if options.make:
            key_manager.make_key(expiration=options.expiration_str)
        else:
            print(key_manager.check_key(key=options.key))
